#include "bitstream.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include "encoder_logger.h"

#define BS_DATA_MAX_LEN (1024 * 1024)

typedef union
{
    uint32_t i;
    uint16_t w[2];
    uint8_t b[4];
} x264_union32_t;

static uint32_t endian_fix32( uint32_t x )
{
    return (x<<24) + ((x<<8)&0xff0000) + ((x>>8)&0xff00) + (x>>24);
}

static uint64_t endian_fix64( uint64_t x )
{
    return endian_fix32(x>>32) + ((uint64_t)endian_fix32(x)<<32);
}

void bs_write(bs_t *s, int n, uint32_t v)
{
    s->cur_bits = (s->cur_bits << n) | v;
    s->i_left -= n;

    if (s->i_left <= 32) {
        ((x264_union32_t*)(s->p))->i = endian_fix64(s->cur_bits << s->i_left);
        s->i_left += 32;
        s->p += 4;
    }
}

static const uint8_t x264_ue_size_tab[256] =
{
     1, 1, 3, 3, 5, 5, 5, 5, 7, 7, 7, 7, 7, 7, 7, 7,
     9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
    11,11,11,11,11,11,11,11,11,11,11,11,11,11,11,11,
    11,11,11,11,11,11,11,11,11,11,11,11,11,11,11,11,
    13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,
    13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,
    13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,
    13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,13,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
    15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,15,
};

/**
 * 有符号哥伦布编码是先对有符号数进行映射，映射为无符号树再进行ue的编码，映射规则如下：
 * v <= 0: -2*v + 1
 * v > 0: 2*v
 */
void bs_write_se( bs_t *s, int val )
{
    int size = 0;
    /* Faster than (val <= 0 ? -val*2+1 : val*2) */
    /* 4 instructions on x86, 3 on ARM */
    int tmp = 1 - val*2;
    if( tmp < 0 ) tmp = val*2;
    val = tmp;

    if( tmp >= 0x100 )
    {
        size = 16;
        tmp >>= 8;
    }
    size += x264_ue_size_tab[tmp];
    bs_write( s, size, val );
}

/**
 * @brief 无符号指数哥伦布编码写入bit流
 * 
 * @param s 
 * @param val 
 */
void bs_write_ue( bs_t *s, unsigned int val )
{
    int size = 0;
    int tmp = ++val;
    if( tmp >= 0x10000 )
    {
        size = 32;
        tmp >>= 16;
    }
    if( tmp >= 0x100 )
    {
        size += 16;
        tmp >>= 8;
    }
    size += x264_ue_size_tab[tmp];
    bs_write( s, size>>1, 0 );
    bs_write( s, (size>>1)+1, val );
}

int bs_init(bs_t *bs)
{
    if (bs == NULL) {
        ERR_LOG("bs is null\n");
        return -1;
    }

    memset(bs, 0, sizeof(bs_t));    
    bs->p_start = (uint8_t *)calloc(BS_DATA_MAX_LEN, 1);
    bs->p = bs->p_start;
    bs->p_end = bs->p_start + BS_DATA_MAX_LEN;
    bs->i_left = sizeof(bs->cur_bits) * 8;
    return 0;
}

void bs_flush( bs_t *s )
{
    (((x264_union32_t*)(s->p))->i) = endian_fix32( s->cur_bits << (s->i_left&31) );
    s->p += 8 - (s->i_left >> 3);
    s->i_left = sizeof(s->cur_bits) * 8;
}

void bs_align_1( bs_t *s )
{
    bs_write( s, s->i_left&7, (1 << (s->i_left&7)) - 1 );
    bs_flush( s );
}